{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "51f516ed-a0dd-452d-b1a4-412d68dde67c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import Dataset\n",
    "from PIL import Image\n",
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7634a0d0-d52d-4e92-b201-554153608d39",
   "metadata": {},
   "outputs": [],
   "source": [
    "建立数据集"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "93f3ad22-84a6-46f5-8349-016eaf98c3b4",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "class mydata(Dataset):\n",
    "    def __init__(self,root_dir,label_dir):\n",
    "        self.root_dir = root_dir\n",
    "        self.label_dir = label_dir\n",
    "        self.path = os.path.join(self.root_dir,self.label_dir)\n",
    "        self.img_path = os.listdir(self.path)\n",
    "    def __getitem__(self,idx):\n",
    "        img_name = self.img_path[idx]\n",
    "        img_item_path = os.path.join(self.path,img_name)\n",
    "        img = Image.open(img_item_path)\n",
    "        label = self.label_dir\n",
    "        return img, label\n",
    "    def __len__(self):\n",
    "        return len(self.img_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "d4b62459-9b2e-4333-8a52-3682eee1d57b",
   "metadata": {},
   "outputs": [],
   "source": [
    "root_dir = \"D:\\\\desktop\\\\practice\\\\数据集\\\\hymenoptera_data\\\\train\"\n",
    "ants_label_dir = \"ants\"\n",
    "ants_dataset = mydata(root_dir,ants_label_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "cde5bef2-9f90-45ff-b578-95c89a90768e",
   "metadata": {},
   "outputs": [],
   "source": [
    "image, label = ants_dataset[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "a9cdb186-2237-4c1e-8a0e-639b2bcd5027",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "124"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "image.show()\n",
    "len(ants_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "0f118d5a-dd5e-4732-87dd-67d527bef548",
   "metadata": {},
   "outputs": [],
   "source": [
    "bees_label_dir = \"bees\"\n",
    "bees_dataset = mydata(root_dir,bees_label_dir)\n",
    "image, label = bees_dataset[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "d4174692-e081-4458-9d4b-ec7224250f11",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "121"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "image.show()\n",
    "len(bees_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e40a741d-2c11-407e-8caa-907bd261156e",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
